package com.zipking.cloud.springbootmybatis.util.excel;
import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.enums.CellDataTypeEnum;
import com.alibaba.excel.metadata.Head;
import com.alibaba.excel.metadata.data.WriteCellData;
import com.alibaba.excel.util.ListUtils;
import com.alibaba.excel.util.MapUtils;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.alibaba.excel.write.metadata.holder.WriteSheetHolder;
import com.alibaba.excel.write.style.column.AbstractColumnWidthStyleStrategy;
import com.zipking.cloud.springbootmybatis.util.JsonUtil;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.ibatis.session.ResultContext;
import org.apache.ibatis.session.ResultHandler;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.poi.ss.usermodel.Cell;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletResponse;
import java.io.BufferedOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URLEncoder;
import java.util.*;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;


public abstract class BaseExportResultHandler<T> implements ResultHandler<T> {

    /**
     * Excel 表头
     */
    private String[] header;

    /**
     * 模板文件名
     */
    private String templateName;

    /**
     * 日志
     */
    private final Logger logger = LogManager.getLogger(BaseExportResultHandler.class.getName());

    /**
     * 模板文件路径
     */
    private final String TEMPLATE_FILE_PATH = "src/main/resources/templates/";

    /**
     * 文件类型
     */
    private final String FILE_TYPE = ".xlsx";

    /**
     * 每隔 100 条数据写入 Excel ，然后清理 list ，方便内存回收
     */
    private final int BATCH_COUNT = 100;

    /**
     * 缓存 list
     */
    private List<T> cachedDataList = ListUtils.newArrayListWithExpectedSize(BATCH_COUNT);

    /**
     * excelWriter
     */
    private ExcelWriter excelWriter;

    /**
     * writeSheet
     */
    private WriteSheet writeSheet;

    /**
     * httpServletResponse
     */
    private HttpServletResponse response;

    /**
     * zip io 输出流
     */
    private ZipOutputStream zos = null;

    /**
     * io 输出流
     */
    private OutputStream os = null;

    /**
     * 是否导出 zip
     */
    private boolean isExportZip;

    /**
     * 重写自适应列宽策略
     */
    static class CustomLongestMatchColumnWidthStyleStrategy extends AbstractColumnWidthStyleStrategy {

        private static final int MAX_COLUMN_WIDTH = 255;

        private final Map<Integer, Map<Integer, Integer>> cache = MapUtils.newHashMapWithExpectedSize(8);

        /**
         * 无参构造
         */
        public CustomLongestMatchColumnWidthStyleStrategy() {
        }

        /**
         * 设定列宽
         *
         * @param writeSheetHolder
         * @param cellDataList
         * @param cell
         * @param head
         * @param relativeRowIndex
         * @param isHead
         */
        @Override
        protected void setColumnWidth(WriteSheetHolder writeSheetHolder, List<WriteCellData<?>> cellDataList, Cell cell, Head head, Integer relativeRowIndex, Boolean isHead) {
            boolean needSetWidth = isHead || !CollectionUtils.isEmpty(cellDataList);
            if (needSetWidth) {
                HashMap<Integer, Integer> maxColumnWidthMap = (HashMap<Integer, Integer>) this.cache.computeIfAbsent(writeSheetHolder.getSheetNo(), (key) -> {
                    return new HashMap<Integer, Integer>(16);
                });
                Integer columnWidth = this.dataLength(cellDataList, cell, isHead);
                if (columnWidth >= 0) {
                    if (columnWidth > MAX_COLUMN_WIDTH) {
                        columnWidth = 255;
                    }

                    Integer maxColumnWidth = (Integer) maxColumnWidthMap.get(cell.getColumnIndex());
                    if (maxColumnWidth == null || columnWidth > maxColumnWidth) {
                        maxColumnWidthMap.put(cell.getColumnIndex(), columnWidth);
                        writeSheetHolder.getSheet().setColumnWidth(cell.getColumnIndex(), columnWidth * 256);
                    }

                }
            }
        }

        private Integer dataLength(List<WriteCellData<?>> cellDataList, Cell cell, Boolean isHead) {
            if (isHead) {
                return cell.getStringCellValue().getBytes().length;
            } else {
                WriteCellData<?> cellData = (WriteCellData) cellDataList.get(0);
                CellDataTypeEnum type = cellData.getType();
                if (type == null) {
                    return -1;
                } else {
                    switch (type) {
                        case STRING:
                            return cellData.getStringValue().getBytes().length;
                        case BOOLEAN:
                            return cellData.getBooleanValue().toString().getBytes().length;
                        case NUMBER:
                            return cellData.getNumberValue().toString().getBytes().length;
                        default:
                            return -1;
                    }
                }
            }
        }
    }

    /**
     * 解析 header 返回
     *
     * @return header 列表
     */
    private List<List<String>> getXlsTemplateHeader(String[] header) {
        List<List<String>> list = new ArrayList<List<String>>();
        for (String h : header) {
            List<String> head = new ArrayList<String>();
            head.add(h);
            list.add(head);
        }
        return list;
    }

    /**
     * 模板内容
     * 1、模板注意用 {} 来表示你要用的变量 如果本来就有"{","}" 特殊字符 用 "\{","\}" 代替
     * 2、填充 list 的时候还要注意 模板中 {.} 多了个点 表示 list
     * 3、如果填充 list 的对象是 map , 必须包涵所有 list 的 key , 哪怕数据为 null，必须使用 map.put(key,null)
     *
     * @param header header
     * @return List<List < Object>>
     */
    private List<List<Object>> getXlsTemplateContent(String[] header) {
        List<List<Object>> list = new ArrayList<>();
        List<Object> templateContent = new ArrayList<>(header.length);
        for (String h : header) {
            templateContent.add(String.format("{.%s}", h));
        }
        list.add(templateContent);
        return list;
    }

    /**
     * 通过 ResultHandler 流式查询数据，由子类自定义实现
     */
    public abstract void fetchDataByStream();

    /**
     * 处理结果，这里可以获取当前行的数据以及行数
     *
     * @param resultContext 上下文
     */
    @Override
    public void handleResult(ResultContext<? extends T> resultContext) {
        T resultObject = resultContext.getResultObject();
        invoke(resultObject);
    }

    /**
     * 处理数据
     *
     * @param obj 泛型对象
     */
    private void invoke(T obj) {
        // 先塞到缓存中
        cachedDataList.add(obj);
        // 如果到达缓存值，则写入 Excel
        if (cachedDataList.size() >= BATCH_COUNT) {
            // 写入 Excel
            writeObjToExcel();
            // 清空缓存数组
            cachedDataList.clear();
        }
    }

    /**
     * 把缓存数据 追加写入 Excel
     */
    private void writeObjToExcel() {
//        logger.info("writeObjToExcel {} ", cachedDataList.size());
        this.excelWriter.fill(cachedDataList, this.writeSheet);
    }

    /**
     * 构造函数
     *
     * @param header Excel 表头
     * @throws IOException io 异常
     */
    public BaseExportResultHandler(String[] header) throws IOException {
        this(header, "DefaultFileName");
    }

    /**
     * 构造函数
     *
     * @param header       Excel 表头
     * @param templateName 模板名
     * @throws IOException io 异常
     */
    public BaseExportResultHandler(String[] header, String templateName) throws IOException {
        this(header, templateName, false);
    }

    /**
     * 构造函数
     *
     * @param header       Excel 表头
     * @param templateName 模板名
     * @param isExportZip  是否导出 zip
     * @throws IOException io 异常
     */
    public BaseExportResultHandler(String[] header, String templateName, boolean isExportZip) throws IOException {
        this.header = header;
        this.templateName = templateName;
        this.isExportZip = isExportZip;
        this.initExcelWriter();
    }

    /**
     * 初始化 Excel 属性
     */
    private void initExcelWriter() throws IOException {
        // http 响应
        //        this.response = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getResponse();
        this.response = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getResponse();
        // io 流
        assert this.response != null;
        this.os = new BufferedOutputStream(this.response.getOutputStream());
        // 如果需要压缩成 zip
        if (this.isExportZip) {
            this.zos = new ZipOutputStream(os);
            ZipEntry zipEntry = new ZipEntry((this.templateName + this.FILE_TYPE).replaceAll(" ", ""));
            this.zos.putNextEntry(zipEntry);
        }
        // 模板路径
        String templateFile = this.TEMPLATE_FILE_PATH + this.templateName + this.FILE_TYPE;
        // 生成模板
        this.generateExcelTemplate(templateFile);
        // 读取模板
        this.excelWriter = EasyExcel
                .write(this.isExportZip ? this.zos : this.os)
                .withTemplate(templateFile)
                .build();
        // 填充模板 sheet
        this.writeSheet = EasyExcel
                .writerSheet()
                .registerWriteHandler(new CustomLongestMatchColumnWidthStyleStrategy())
                .build();
    }

    /**
     * 生成模板
     *
     * @param templateFile 模板名
     */
    private void generateExcelTemplate(String templateFile) {
        EasyExcel
                .write(templateFile)
                //.registerWriteHandler(getCustomHorizontalCellStyleStrategy())
                .registerWriteHandler(new CustomLongestMatchColumnWidthStyleStrategy())
                .head(this.getXlsTemplateHeader(this.header))
                .sheet(this.templateName)
                .doWrite(this.getXlsTemplateContent(this.header));
    }

    /**
     * 开始导出数据
     *
     * @throws IOException io 异常
     */
    public void startExportExcel()   {

        try {
            this.response.setCharacterEncoding("utf-8");
            // 下载文件名，这里URLEncoder.encode可以防止中文乱码
            String fileName = URLEncoder.encode(this.templateName, "UTF-8").replaceAll("\\+", "%20");
            if (!isExportZip) {
                // 这里注意 使用swagger 会导致各种问题，请直接用浏览器或者用postman
                this.response.setContentType("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet");
                this.response.setHeader("Content-disposition", "attachment;filename*=utf-8''" + fileName + ".xlsx");
                logger.info("is Export Excel");
            } else {
                this.response.setContentType("application/octet-stream");
                this.response.setHeader("Content-disposition", "attachment;filename*=utf-8''" + fileName + ".zip");
                logger.info("is Export Zip");
            }
            // 子类流式获取数据
            this.fetchDataByStream();
        } catch (Exception e) {
            // 重置response
            response.reset();
            response.setContentType("application/json");
            response.setCharacterEncoding("utf-8");
            Map<String, String> map = MapUtils.newHashMap();
            map.put("status", "failure");
            map.put("message", "下载文件失败" + e.getMessage());
            try {
                response.getWriter().println(JsonUtil.toStr(map));
            } catch (IOException ioException) {
                ioException.printStackTrace();
            }
        } finally {
            // 关闭 Excel 资源
            this.excelWriter.finish();
            // 关闭 io 资源
            if (this.isExportZip) {
                try {
                    if (zos != null) {
                        zos.close();
                    }
                } catch (IOException e) {
                    logger.error("error", e);
                }
            } else {
                try {
                    if (os != null) {
                        os.close();
                    }
                } catch (IOException e) {
                    logger.error("error", e);
                }
            }
        }
    }
}

